import torch.cuda
from torch import nn

from common_models.MLP import MLP
from common_utils.losses import hinge_loss, svm_loss

device = 'cuda' if torch.cuda.is_available() else 'cpu'
import torch.nn.functional as F

class AttrModel(nn.Module):
    """
     每个属性一个分类器，
     损失函数使用BCELoss
    """
    def __init__(self, args,feat_dim,attr_dim):
        super(AttrModel, self).__init__()
        self.attr_classifier_list = []
        # 输入维度为图片特征维度，输出维度为属性的数量
        for idx in range(attr_dim):
            attr_classifier = MLP(feat_dim, 1, relu=args.dropout,dropout=args.dropout,norm=args.norm).to(device)
            self.attr_classifier_list.append(attr_classifier)
        # self.img_classifier = MLP(attr_dim, len(dataset.classes), relu=False).to(device)

    def forward(self, imgs, labels=None, attrs=None):
        """
        .. math:: MSELoss： \mathrm{loss}(\mathbf{x}_i,\mathbf{y}_i)=(\mathbf{x}_i-\mathbf{y}_i)^2
        :param imgs:
        :param labels:
        :param attrs:
        :return:
        """
        attr_preds = torch.Tensor().to(device)
        # loss_fn = torch.nn.BCELoss()
        loss_fn = torch.nn.MSELoss()
        loss = []
        # crossEntropyLoss = nn.CrossEntropyLoss()
        for idx,attr_classifier in enumerate(self.attr_classifier_list):
            # 根据图片特征学习对象的属性
            attr_pred = F.tanh(attr_classifier(imgs))
            # attr_pred = attr_classifier(imgs)
            attr_item = attrs[:,idx].view(len(imgs),1)
            loss_item = loss_fn(attr_pred, attr_item)
            loss.append(loss_item)
            # attr_pred2 = torch.max(attr_pred,dim=1)[1].unsqueeze(1)
            attr_preds = torch.cat((attr_preds,attr_pred),dim=1)
        # label_preds = self.img_classifier(attr_preds)
        # loss_label = crossEntropyLoss(label_preds,labels)
        return loss,attr_preds,None
